#!/bin/bash

# SD3训练脚本 - 16GB显存版本
# 基于官方文档配置

# 基础路径定义
WORKSPACE="/workspace"
AI_DIR="$WORKSPACE/AI"
PYTHON_PATH="$AI_DIR/.train/bin/python"
SCRIPT_PATH="$AI_DIR/ComfyUI/custom_nodes/comfyui_lora_train/sd_scripts/sd3_train_network.py"

# 模型路径定义
MODELS_DIR="$AI_DIR/ComfyUI/models"
SD3_MODEL="$MODELS_DIR/checkpoints/sd3.5_large.safetensors"
CLIP_L_MODEL="$MODELS_DIR/clip/clip_l.safetensors"
CLIP_G_MODEL="$MODELS_DIR/clip/clip_g.safetensors"
T5XXL_MODEL="$MODELS_DIR/text_encoders/t5xxl_fp16.safetensors"

# 数据集和输出路径
TRAIN_DATA_DIR="$WORKSPACE/training_data/images"
OUTPUT_DIR="$WORKSPACE/training_data/output"

# 设置环境变量以优化内存分配
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

# 16GB显存优化参数
$PYTHON_PATH $SCRIPT_PATH \
    --pretrained_model_name_or_path "$SD3_MODEL" \
    --clip_l "$CLIP_L_MODEL" \
    --clip_g "$CLIP_G_MODEL" \
    --t5xxl "$T5XXL_MODEL" \
    --cache_latents_to_disk \
    --save_model_as safetensors \
    --sdpa \
    --persistent_data_loader_workers \
    --max_data_loader_n_workers 1 \
    --seed 42 \
    --gradient_checkpointing \
    --mixed_precision bf16 \
    --save_precision bf16 \
    --network_module networks.lora_sd3 \
    --network_dim 2 \
    --network_train_unet_only \
    --optimizer_type Adafactor \
    --learning_rate 1e-4 \
    --cache_text_encoder_outputs \
    --cache_text_encoder_outputs_to_disk \
    --fp8_base \
    --highvram \
    --max_train_steps 200 \
    --save_every_n_epochs 1 \
    --train_data_dir "$TRAIN_DATA_DIR" \
    --output_dir "$OUTPUT_DIR" \
    --output_name sd3-lora-16gb \
    --resolution 512 \
    --train_batch_size 1 \
    --optimizer_args "relative_step=False" "scale_parameter=False" "warmup_init=False" \
    --lr_scheduler constant_with_warmup \
    --max_grad_norm 0.0 \
    --disable_mmap_load_safetensors \
    --enable_bucket \
    --min_bucket_reso 256 \
    --max_bucket_reso 512 \
    --bucket_reso_steps 64

echo "SD3训练完成！输出保存在: $OUTPUT_DIR" 